package com.example.hook.agent;

import javassist.*;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;

/**
 * Hook Agent - 专门用来Hook外部应用程序
 */
public class HookAgent {
    
    public static void premain(String agentArgs, Instrumentation inst) {
        System.out.println("=== Hook Agent 启动 ===");
        System.out.println("Agent参数: " + agentArgs);
        System.out.println("准备Hook目标应用程序...");
        
        // 注册类转换器
        inst.addTransformer(new TargetAppTransformer(), true);
        
        System.out.println("Hook Agent 初始化完成");
    }
    
    public static void agentmain(String agentArgs, Instrumentation inst) {
        System.out.println("=== Hook Agent 动态附加 ===");
        premain(agentArgs, inst);
    }
    
    /**
     * 目标应用程序类转换器
     */
    static class TargetAppTransformer implements ClassFileTransformer {
        
        @Override
        public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
                              ProtectionDomain protectionDomain, byte[] classfileBuffer) {
            
            // 只Hook目标应用程序的类
            if (className != null && className.equals("com/example/target/TargetService")) {
                System.out.println("正在Hook目标应用类: " + className);
                return transformTargetService(classfileBuffer);
            }
            
            return null; // 不修改其他类
        }
        
        /**
         * 转换目标服务类
         */
        private byte[] transformTargetService(byte[] classfileBuffer) {
            try {
                ClassPool pool = ClassPool.getDefault();
                CtClass cc = pool.makeClass(new java.io.ByteArrayInputStream(classfileBuffer));
                
                // Hook add方法
                hookAddMethod(cc);
                
                // Hook getUserInfo方法
                hookGetUserInfoMethod(cc);
                
                // Hook formatMessage静态方法
                hookFormatMessageMethod(cc);
                
                // Hook processOrder方法
                hookProcessOrderMethod(cc);
                
                // Hook queryDatabase方法
                hookQueryDatabaseMethod(cc);
                
                byte[] byteCode = cc.toBytecode();
                cc.detach(); // 释放内存
                
                System.out.println("目标应用类Hook完成: com.example.target.TargetService");
                return byteCode;
                
            } catch (Exception e) {
                System.err.println("Hook目标应用类失败: " + e.getMessage());
                e.printStackTrace();
                return null;
            }
        }
        
        /**
         * Hook add方法 - 将结果乘以10
         */
        private void hookAddMethod(CtClass cc) throws Exception {
            CtMethod addMethod = cc.getDeclaredMethod("add");
            
            addMethod.insertBefore("System.out.println(\"[Hook Agent] 拦截add方法，参数: \" + $1 + \", \" + $2);");
            addMethod.insertAfter("{ " +
                "System.out.println(\"[Hook Agent] add方法原返回值: \" + $_ + \", 修改为: \" + ($_ * 10)); " +
                "$_ = $_ * 10; " +
                "}");
        }
        
        /**
         * Hook getUserInfo方法 - 特殊用户ID返回VIP信息
         */
        private void hookGetUserInfoMethod(CtClass cc) throws Exception {
            CtMethod getUserInfoMethod = cc.getDeclaredMethod("getUserInfo");
            
            getUserInfoMethod.insertBefore("System.out.println(\"[Hook Agent] 拦截getUserInfo方法，用户ID: \" + $1);");
            getUserInfoMethod.insertAfter("{ " +
                "System.out.println(\"[Hook Agent] getUserInfo原返回值: \" + $_); " +
                "if ($1 == 999) { " +
                "    $_ = \"VIP_GUEST_999\"; " +
                "    System.out.println(\"[Hook Agent] 特殊用户999，返回值修改为: \" + $_); " +
                "} else if ($1 == 1) { " +
                "    $_ = \"SUPER_ADMIN\"; " +
                "    System.out.println(\"[Hook Agent] 管理员用户1，返回值修改为: \" + $_); " +
                "}" +
                "}");
        }
        
        /**
         * Hook formatMessage静态方法 - 添加Hook标识
         */
        private void hookFormatMessageMethod(CtClass cc) throws Exception {
            CtMethod formatMessageMethod = cc.getDeclaredMethod("formatMessage");
            
            formatMessageMethod.insertBefore("System.out.println(\"[Hook Agent] 拦截静态方法formatMessage，消息: \" + $1);");
            formatMessageMethod.insertAfter("{ " +
                "System.out.println(\"[Hook Agent] formatMessage原返回值: \" + $_); " +
                "$_ = \"[HOOKED] \" + $_; " +
                "System.out.println(\"[Hook Agent] formatMessage修改后返回值: \" + $_); " +
                "}");
        }
        
        /**
         * Hook processOrder方法 - 添加订单验证和修改
         */
        private void hookProcessOrderMethod(CtClass cc) throws Exception {
            CtMethod processOrderMethod = cc.getDeclaredMethod("processOrder");
            
            processOrderMethod.insertBefore("{ " +
                "System.out.println(\"[Hook Agent] 拦截processOrder方法，订单ID: \" + $1 + \", 金额: \" + $2); " +
                "if ($2 > 2000.0) { " +
                "    System.out.println(\"[Hook Agent] 检测到高金额订单，需要特殊处理\"); " +
                "}" +
                "}");
            
            processOrderMethod.insertAfter("{ " +
                "System.out.println(\"[Hook Agent] processOrder原返回值: \" + $_); " +
                "if ($2 > 1000.0 && $2 <= 2000.0) { " +
                "    $_ = \"PREMIUM_\" + $_; " +
                "    System.out.println(\"[Hook Agent] 高价值订单，返回值修改为: \" + $_); " +
                "} else if ($2 > 2000.0) { " +
                "    $_ = \"ULTRA_PREMIUM_\" + $_; " +
                "    System.out.println(\"[Hook Agent] 超高价值订单，返回值修改为: \" + $_); " +
                "}" +
                "}");
        }
        
        /**
         * Hook queryDatabase方法 - 添加性能监控和缓存模拟
         */
        private void hookQueryDatabaseMethod(CtClass cc) throws Exception {
            CtMethod queryDatabaseMethod = cc.getDeclaredMethod("queryDatabase");
            
            queryDatabaseMethod.insertBefore("{ " +
                "System.out.println(\"[Hook Agent] 拦截数据库查询，表: \" + $1 + \", 条件: \" + $2); " +
                "long startTime = System.currentTimeMillis(); " +
                "com.example.hook.agent.HookAgent.setQueryStartTime(startTime); " +
                "}");
            
            queryDatabaseMethod.insertAfter("{ " +
                "long endTime = System.currentTimeMillis(); " +
                "long duration = endTime - com.example.hook.agent.HookAgent.getQueryStartTime(); " +
                "System.out.println(\"[Hook Agent] 数据库查询耗时: \" + duration + \"ms\"); " +
                "System.out.println(\"[Hook Agent] 查询原返回值: \" + $_); " +
                "if (duration > 80) { " +
                "    $_ = \"[SLOW_QUERY] \" + $_; " +
                "    System.out.println(\"[Hook Agent] 检测到慢查询，标记返回值: \" + $_); " +
                "} else { " +
                "    $_ = \"[CACHED] \" + $_; " +
                "    System.out.println(\"[Hook Agent] 快速查询，添加缓存标记: \" + $_); " +
                "}" +
                "}");
        }
    }
    
    // 静态变量用于存储查询开始时间
    private static long queryStartTime = 0;
    
    public static void setQueryStartTime(long time) {
        queryStartTime = time;
    }
    
    public static long getQueryStartTime() {
        return queryStartTime;
    }
}